from turtle import forward
import torch
import torch.nn as nn
import random
from data_loader.transforms import mask
from data_loader.transforms.mask import compute_mask_indices, get_mask

class RandomAugment(nn.Module):
    def __init__(self, aug_intensity: int, label_pool: int):
        super().__init__()
        assert aug_intensity < 11, "augmentation intensity range [1, 10]"
        self.aug_intensity = aug_intensity
        self.label_pool = label_pool
    
    @torch.no_grad()
    def forward(self, spec, labels):
        if self.aug_intensity == 0:
            return spec, labels
        else:
            spec = spec.T
            labels = labels.T
        aug_type = random.randint(0, 4)
        if aug_type == 0:
            # No augmentation
            pass
        elif aug_type == 1:
            spec, labels = self.frame_shift(spec, labels)
        elif aug_type == 2:
            # spec, labels = self.time_mask(spec, labels)
            pass
        elif aug_type == 3:
            # spec, labels = self.freq_mask(spec, labels)
            pass
        elif aug_type == 4:
            spec, labels = self.filter_aug(spec, labels)
        return spec.T, labels.T
    
    
    def frame_shift(self, spec, labels):
        T, C = labels.shape
        shift_scale = random.randint(1, self.aug_intensity) / 10
        shift_len_label = int(T * shift_scale)
        shift_len_spec = int(shift_len_label * self.label_pool)
        
        spec = torch.roll(spec, shifts=shift_len_spec, dims=0)
        labels = torch.roll(labels, shifts=shift_len_label, dims=0)
               
        return spec, labels

    def time_mask(self, spec, labels):
        T, C = labels.shape
        mask_labels = compute_mask_indices(
            shape=(1, T),
            mask_prob=self.aug_intensity / 10,
            mask_length=4,
            no_overlap=True,
            min_masks=2,
            mask_type="static",
            padding_mask=None
        )
        mask_labels = torch.arange(0, T)[mask_labels[0]]
        mask_spec = (mask_labels * self.label_pool).unsqueeze(1)
        mask_spec = mask_spec + torch.arange(0, self.label_pool).unsqueeze(0)   # broadcast
        mask_spec = mask_spec.flatten()
        
        masked_spec = spec.clone()
        masked_labels = labels.clone()
        masked_spec[mask_spec] *= 0
        masked_labels[mask_labels] *= 0
        
        return masked_spec, masked_labels
    
    def freq_mask(self, spec, labels):
        T, F = spec.shape
        mask_bands = compute_mask_indices(
            shape=(T, F),
            mask_prob=self.aug_intensity / 10,
            mask_length=4,
            no_overlap=True,
            min_masks=2,
            mask_type="static",
            padding_mask=None
        )
        mask_bands = torch.tensor(mask_bands, dtype=torch.float)
        return (1 - mask_bands) * spec, labels
    
    def filter_aug(self, spec, labels):
        T, F = spec.shape
        db_range = (-7.5, 6)
        n_freq_band = torch.randint(low=2, high=self.aug_intensity, size=(1,)).item()  # [low, high)
        band_bndry_freqs = torch.cat((
            torch.tensor([0]),
            torch.sort(torch.randint(1, F - 1, (n_freq_band - 1,)))[0],
            torch.tensor([F])))
        band_factors = torch.rand((1, n_freq_band)) * (db_range[1] - db_range[0]) + \
                        db_range[0]
        band_factors = 10 ** (band_factors / 20)
        freq_filt = torch.ones((1, F))
        for i in range(n_freq_band):
            freq_filt[:, band_bndry_freqs[i]:band_bndry_freqs[i + 1]] = band_factors[:, i].unsqueeze(-1).unsqueeze(-1)
        return spec * freq_filt, labels


if __name__ == "__main__":
    # Example usage
    torch.manual_seed(42)  # For reproducibility
    random.seed(4)
    spec = torch.randn(626, 128) * torch.arange(626).unsqueeze(1)  # Example spectrogram with 100 time steps and 80 frequency bins
    labels = torch.rand(156, 10)  # Example labels with 10 classes
    
    augmenter = RandomAugment(aug_intensity=6, label_pool=4)
    for i in range(100):
        augmented_spec, augmented_labels = augmenter(spec, labels)
    
    
    # print("Original Spec Shape:", spec.shape)
    # print("Augmented Spec Shape:", augmented_spec.shape)
    # print("Original Labels Shape:", labels.shape)
    # print("Augmented Labels Shape:", augmented_labels.shape)
    
    # # plot the original and augmented spectrograms
    # import matplotlib.pyplot as plt
    # plt.figure(figsize=(12, 6))
    # plt.subplot(2, 1, 1)
    # plt.imshow(spec.T.numpy(), aspect='auto', origin='lower')
    # plt.title("Original Spectrogram")
    # plt.colorbar()
    # plt.subplot(2, 1, 2)
    # plt.imshow(augmented_spec.T.numpy(), aspect='auto', origin='lower')
    # plt.title("Augmented Spectrogram")
    # plt.colorbar()
    # plt.savefig("spectrogram_augmentation.png")